import datetime
import signal
import tempfile
import os
import subprocess
from absl import app, flags
import jinja2
from tornado.ioloop import IOLoop
from itertools import zip_longest
from coin.tool.strat_monitor.app.monitor_runner_slurm.monitor_runner_slurm import SBatchRunner
from coin.exchange.base.strategy_logger import parse_file_name
from coin.support.proto_log.logic.util import enumerate_strat_log_into_stratinfos


FLAGS = flags.FLAGS


def _grouper(n, iterable, fillvalue=None):
  args = [iter(iterable)] * n
  return zip_longest(fillvalue=fillvalue, *args)


def _get_every_n_strategies(n, strat_infos):
  # TODO(daniel):
  strategies = []
  for strat_info in strat_infos:
    if strat_info.strategy_name not in strategies:
      strategies.append(strat_info.strategy_name)
  strategy_list = []
  for n0, n1, n2, n3, n4, n5, n6, n7, n8, n9 in _grouper(n, strategies):
    tmp = [n0, n1, n2, n3, n4, n5, n6, n7, n8, n9]
    strats = [i for i in tmp if i]
    names = ','.join(strats)
    info = {
        'name': names,
        'mem-per-cpu': '%dM' % (0),  # TODO(daniel):
    }
    strategy_list.append(info)
  return strategy_list


def _get_yesterday():
  return (datetime.date.today() - datetime.timedelta(days=1)).strftime('%Y%m%d')


def generate_summary_report(tmp_dir, trading_date, log_type):
  cmd = f"./pyrunner python/coin/support/proto_log/app/strat_state_anaylyzer.py" \
        f" --generate_summary_report_only" \
        f" --tmp_dir={tmp_dir}" \
        f" --trading_date={trading_date}"
  ret = subprocess.run(cmd, shell=True, encoding="utf-8")
  status = 'fail'
  if ret.returncode == 0:
    status = 'success'
  print("generate summary report:", status)


def main(argv):
  # check dir of strat proto log
  strat_proto_log_dir = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'
  if not os.path.exists(strat_proto_log_dir):
    print('ERROR: dir non-exist %s' % strat_proto_log_dir)
    return
  # check trading date
  trading_date = FLAGS.trading_date
  if trading_date is None:
    trading_date = _get_yesterday()
  # collect strategy name
  strat_infos = enumerate_strat_log_into_stratinfos(strat_proto_log_dir,
                                                    trading_date,
                                                    parse_file_name)
  strategy_list = _get_every_n_strategies(10, strat_infos)
  # prepare srun scripts
  template_loader = jinja2.FileSystemLoader(searchpath='./')
  template_env = jinja2.Environment(loader=template_loader)
  template = template_env.get_template(
      'python/coin/support/proto_log/app/strat_state_anaylyzer_script.tmpl')
  script = template.render(strategy_list=strategy_list,
                           trading_date=trading_date,
                           remote_home=FLAGS.remote_home,
                           coin_env=FLAGS.coin_env,
                           tmp_dir=FLAGS.tmp_dir,
                           log_type=FLAGS.log_type)
  # dry run
  if FLAGS.dry_run:
    print(script)
    return
  # sbatch run
  with tempfile.NamedTemporaryFile(mode='w', dir='') as script_file:
    script_file.write(script)
    script_file.flush()

    ioloop = IOLoop.current()
    sbatch_runner = SBatchRunner(ioloop, script_file.name, FLAGS.sbatch_args)
    signal.signal(signal.SIGTERM, sbatch_runner.signal_handler)
    signal.signal(signal.SIGINT, sbatch_runner.signal_handler)
    job_id = sbatch_runner.run()
    print(job_id)
  # generate summary report
  generate_summary_report(FLAGS.tmp_dir, trading_date, FLAGS.log_type)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('tmp_dir', '/tmp', 'dir for keep tmp report')
  flags.DEFINE_string('remote_home', None, 'path of remote home.')
  flags.DEFINE_string('coin_env',
                      'miniconda3/envs/coin2_motion_env',
                      'coin_env path relative to remote_home')
  flags.DEFINE_bool('dry_run', False, 'print sbatch script only and not run sbatch')
  flags.DEFINE_string('sbatch_args', None, 'sbatch arguments')
  flags.DEFINE_string('log_type', None, '')
  app.run(main)
